package main

import (
	"bufio"
	"encoding/binary"
	"errors"
	"io"
	"net"
)

// 解析握手包
func UnpackHandshakePacket(conn net.Conn) HandshakePacket {
	var packet HandshakePacket
	buffer := make([]byte, 2)
	reader := bufio.NewReader(conn)

	_, err := io.ReadFull(reader, buffer)
	if err != nil {
		packet.Err = err
		return packet
	}
	packet.Ver, packet.NMethods = buffer[0], buffer[1]
	packet.Methods = make([]byte, packet.NMethods)
	_, err = io.ReadFull(reader, packet.Methods)
	if err != nil {
		packet.Err = err
	}

	return packet
}

// 解析请求包
func UnpackRequestPacket(conn net.Conn) RequestPacket {
	var packet RequestPacket
	buffer := make([]byte, 4)
	reader := bufio.NewReader(conn)

	_, err := io.ReadFull(reader, buffer)
	if err != nil {
		packet.Err = err
		return packet
	}
	packet.Ver, packet.Cmd, packet.Rsv, packet.Atyp = buffer[0], buffer[1], buffer[2], buffer[3]

	if packet.Atyp == SOCKS5_ADDR_TYPE_IPV4 {
		packet.DstAddr = make([]byte, 4)
		_, err := io.ReadFull(reader, packet.DstAddr)
		if err != nil {
			packet.Err = err
			return packet
		}
	} else if packet.Atyp == SOCKS5_ADDR_TYPE_IPV6 {
		packet.DstAddr = make([]byte, 16)
		_, err := io.ReadFull(reader, packet.DstAddr)
		if err != nil {
			packet.Err = err
			return packet
		}
	} else if packet.Atyp == SOCKS5_ADDR_TYPE_DOMAINNAME {
		size, err := reader.ReadByte()
		if err != nil {
			packet.Err = err
			return packet
		}
		packet.DstAddr = make([]byte, size)
		_, err = io.ReadFull(reader, packet.DstAddr)
		if err != nil {
			packet.Err = err
			return packet
		}
	} else {
		packet.Err = errors.New("unknown address type")
	}

	buffer = buffer[:0]
	_, err = io.ReadFull(reader, buffer[:2])
	if err != nil {
		packet.Err = err
		return packet
	}
	packet.DstPort = binary.BigEndian.Uint16(buffer[:2])

	return packet
}
